#!/usr/bin/python

##
## Copyright (C) 2009 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

from expression import *
from knowngene import *
import subprocess

import time
import sys

try:
    import pysqlite2.dbapi2 as sqlite
except:
    import sqlite3 as sqlite

def printLocation(kg, chromosome, start, stop):
       gene=[]
       whats='INTERGENE'
       try:
           istart=int(start)
           istop=int(stop)

           gene=kg.getCDSGeneCached(chromosome, istart, istop)
           if len(gene)>0:
               whats='CDS'

           if len(gene)==0:
               gene=kg.getUTRGeneCached(chromosome, istart, istop)
               if len(gene)>0:
                   whats='UTR'

           if len(gene)==0:
               gene=kg.getExonGeneCached(chromosome, istart, istop)
               if len(gene)>0:
                   whats='EXON'

           if len(gene)==0:
               gene=kg.getGeneGeneCached(chromosome, istart, istop)
               if len(gene)>0:
                   whats='GENE'

       except:
           gene=[]

       if len(gene)==0:
           gene=[ ['='] ]

       row=[]
       for gg in gene:
           for tt in kg.getAnnotationsCached(gg[0]):
               if tt not in row:
                   row.append( tt )
       return [ ",".join(row), whats ]


###########################################################################

def main():
   """main routine"""

   # pars
   database=''
   dbsnp=''
   filesnp=''

   errorMsg="SYNTAX: " + sys.argv[0] + " <DataBase> <snpDB.sqlite> <fileMAQ.snp>\n"

   if len(sys.argv)<2:
       print >>sys.stderr, errorMsg
       return

   for a in sys.argv[1:]:
       if database=='':
           database=a
           continue
       if dbsnp=='':
           dbsnp=a
           continue
       if filesnp=='':
           filesnp=a

   if database=='' or filesnp=='':
       print >>sys.stderr, errorMsg
       return

   print >>sys.stderr, time.ctime() + " CACHE ANNOTATION DATA"

   kg=knowngene(database)
   kg.createCache()

   sql=sqlite.connect(dbsnp)
   sql.execute("PRAGMA CACHE_SIZE = 5000000")
   cur = sql.cursor()

   print >>sys.stderr, time.ctime() + " START MAP"

   fout = open(filesnp, "a")
   for line in open(filesnp, "r"):
       line=line.strip()
       process=re.split('\s', line)
       chromosome=process[0]
       if len(chromosome)<4 or chromosome[0:3]!="chr":
           chromosome="chr"+chromosome
       snpstop=int(process[1])
       snpstart=snpstop-1
       # dbsnp have position with 0 start and maq has 1 start
       gene=printLocation(kg, chromosome, snpstart, snpstop)
       # locate SNP if annotated
       snpdata=[ line, gene[0], gene[1], '=', '=', '=' ]
       cur.execute("select * from snp where chromosome=? and sstop=?;", (chromosome, snpstop) )
       for e in cur:
           snpdata=[line, gene[0], gene[1], e[2], e[3], e[4] ]
       print "\t".join( snpdata )

   print >>sys.stderr, time.ctime() + " END MAP"

if __name__=='__main__': main()
